import threading
import selectors
import socket
import time
from concurrent.futures import Future, ThreadPoolExecutor
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("SimpleAsyncIO")


class SimpleAsyncIO:
    def __init__(self, max_workers=10):
        self.selector = selectors.DefaultSelector()
        self.executor = ThreadPoolExecutor(max_workers=max_workers)
        self.loop_thread = None
        self.running = False
        self.futures = {}          # socket -> Future

    # ---------- 以下 4 个方法提交给线程池 ----------
    def _do_read(self, sock, future):
        """真正的 read 逻辑，跑在线程池里"""
        try:
            data = sock.recv(future.buffer_size)
            if data:
                future.set_result(data)
            else:
                future.set_exception(ConnectionError("Connection closed"))
        except Exception as e:
            future.set_exception(e)

    def _do_write(self, sock, future):
        """真正的 write 逻辑，跑在线程池里"""
        try:
            total_sent = 0
            while total_sent < len(future.data):
                sent = sock.send(future.data[total_sent:])
                if sent == 0:
                    raise ConnectionError("Socket send zero bytes")
                total_sent += sent
            future.set_result(total_sent)
        except Exception as e:
            future.set_exception(e)

    def _do_accept(self, sock, future):
        """真正的 accept 逻辑，跑在线程池里"""
        try:
            client_sock, addr = sock.accept()
            client_sock.setblocking(False)
            future.set_result((client_sock, addr))
        except Exception as e:
            future.set_exception(e)

    def _do_connect(self, sock, future):
        """真正的 connect 逻辑，跑在线程池里"""
        try:
            error_code = sock.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
            if error_code == 0:
                future.set_result(None)
            else:
                future.set_exception(ConnectionError(f"Connect error {error_code}"))
        except Exception as e:
            future.set_exception(e)
    # ------------------------------------------------

    # ===================== 以下代码基本不变 =====================
    def start(self):
        if self.running:
            return
        self.running = True
        self.loop_thread = threading.Thread(target=self._event_loop)
        self.loop_thread.daemon = True
        self.loop_thread.start()
        logger.info("Event loop started")

    def stop(self):
        self.running = False
        if self.loop_thread:
            self.loop_thread.join(timeout=1.0)
        self.executor.shutdown(wait=False)
        logger.info("Event loop stopped")

    def _event_loop(self):
        """事件循环：检测到事件后，把任务提交给线程池"""
        while self.running:
            events = self.selector.select(timeout=0.1)
            for key, mask in events:
                sock = key.fileobj
                future = self.futures.get(sock)
                if not future:
                    self.selector.unregister(sock)
                    continue

                # 只提交一次，提交后立刻取消注册，防止重复提交
                if mask & selectors.EVENT_READ and future.type == 'read':
                    self.selector.unregister(sock)
                    del self.futures[sock]
                    self.executor.submit(self._do_read, sock, future)

                elif mask & selectors.EVENT_WRITE and future.type == 'write':
                    self.selector.unregister(sock)
                    del self.futures[sock]
                    self.executor.submit(self._do_write, sock, future)

                elif mask & selectors.EVENT_READ and future.type == 'accept':
                    # accept 不取消注册，服务器需要持续 accept
                    # 为了避免重复提交，这里用 Future 的 state 做简单保护
                    if not future.set_running_or_notify_cancel():
                        continue
                    self.executor.submit(self._do_accept, sock, future)

                elif mask & selectors.EVENT_WRITE and future.type == 'connect':
                    self.selector.unregister(sock)
                    del self.futures[sock]
                    self.executor.submit(self._do_connect, sock, future)

    # ===================== 对外 API 不变 =====================
    def read(self, sock, size):
        future = ReadFuture(size)
        self.futures[sock] = future
        self.selector.register(sock, selectors.EVENT_READ, future)
        return future

    def write(self, sock, data):
        future = WriteFuture(data)
        self.futures[sock] = future
        self.selector.register(sock, selectors.EVENT_WRITE, future)
        return future

    def accept(self, sock):
        future = AcceptFuture()
        self.futures[sock] = future
        self.selector.register(sock, selectors.EVENT_READ, future)
        return future

    def connect(self, sock, address):
        future = ConnectFuture()
        self.futures[sock] = future
        try:
            sock.connect(address)
            future.set_result(None)
            return future
        except BlockingIOError:
            self.selector.register(sock, selectors.EVENT_WRITE, future)
            return future
        except Exception as e:
            future.set_exception(e)
            return future


# ------------- 以下代码完全不改动 -------------
class ReadFuture(Future):
    def __init__(self, buffer_size):
        super().__init__()
        self.type = 'read'
        self.buffer_size = buffer_size

class WriteFuture(Future):
    def __init__(self, data):
        super().__init__()
        self.type = 'write'
        self.data = data
        self.offset = 0

class AcceptFuture(Future):
    def __init__(self):
        super().__init__()
        self.type = 'accept'

class ConnectFuture(Future):
    def __init__(self):
        super().__init__()
        self.type = 'connect'

class AsyncSocket:
    def __init__(self, sock=None, async_io=None):
        self.sock = sock or socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.async_io = async_io
        self.sock.setblocking(False)

    @classmethod
    def create(cls, async_io, family=socket.AF_INET, type=socket.SOCK_STREAM):
        sock = socket.socket(family, type)
        return cls(sock, async_io)

    def bind(self, address):
        self.sock.bind(address)
        return self

    def listen(self, backlog=5):
        self.sock.listen(backlog)
        return self

    def accept(self):
        if not self.async_io:
            raise ValueError("AsyncIO instance not provided")
        return self.async_io.accept(self.sock)

    def connect(self, address):
        if not self.async_io:
            raise ValueError("AsyncIO instance not provided")
        return self.async_io.connect(self.sock, address)

    def read(self, size):
        if not self.async_io:
            raise ValueError("AsyncIO instance not provided")
        return self.async_io.read(self.sock, size)

    def write(self, data):
        if not self.async_io:
            raise ValueError("AsyncIO instance not provided")
        return self.async_io.write(self.sock, data)

    def close(self):
        if self.sock:
            self.sock.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


# -------------- 示例代码 --------------
MESSAGE_LENGTH = 200
NITERS = 1000

def create_fixed_length_message(content):
    content = content[:MESSAGE_LENGTH].ljust(MESSAGE_LENGTH)
    return content.encode()

def server_example():
    async_io = SimpleAsyncIO()
    async_io.start()
    try:
        server_socket = AsyncSocket.create(async_io)
        server_socket.bind(('localhost', 8888))
        server_socket.listen()
        print("Server started, waiting for connections...")
        accept_future = server_socket.accept()
        client_sock, addr = accept_future.result()
        print(f"Client connected: {addr}")
        client_async_socket = AsyncSocket(client_sock, async_io)
        for i in range(NITERS):
            read_future = client_async_socket.read(MESSAGE_LENGTH)
            client_message = read_future.result().decode().strip()
            print(f"Received from client: {client_message}")
            response = create_fixed_length_message(f"Hello from server {i}th iter!")
            write_future = client_async_socket.write(response)
            write_future.result()
            print("Response sent to client")
    except Exception as e:
        print(f"Server error: {e}")
    finally:
        async_io.stop()

def client_example():
    async_io = SimpleAsyncIO()
    async_io.start()
    try:
        client_socket = AsyncSocket.create(async_io)
        connect_future = client_socket.connect(('localhost', 8888))
        connect_future.result()
        print("Connected to server")
        
        for i in range(NITERS):
            message = create_fixed_length_message(f"Hello from client {i}th iter!") 
            write_future = client_socket.write(message)
            write_future.result()
            print("Message sent to server")
            read_future = client_socket.read(MESSAGE_LENGTH)
            server_response = read_future.result().decode().strip()
            print(f"Received from server: {server_response}")
    except Exception as e:
        print(f"Client error: {e}")
    finally:
        async_io.stop()

if __name__ == "__main__":
    import sys
    if sys.argv[1] == 's':
        server_example()
    else:
        client_example()